from Dataloader.weiboloader import *
from SentModel.Sent2Vec import *
from RumdetecFramework.GraphAutoEncoder import *
import torch.nn as nn

tr = WeiboGraphSet(max_seq_len=10000, min_seq_len=5, DataAug=True)
tr.load_data_fast(data_prefix="../data/WeiboGraph_tr")

te = WeiboGraphSet(max_seq_len=10000, min_seq_len=5)
te.load_data_fast(data_prefix="../data/WeiboGraph_te")

dev = WeiboGraphSet(max_seq_len=10000, min_seq_len=5)
dev.load_data_fast(data_prefix="../data/WeiboGraph_dev")

w2v = W2V_SelfAttn("../saved/word2vec_cn/", 300)
gat = GAT(300, 300)
GAE_Trainer = GraphAutoEncoderTrainer(w2v, gat, batch_size=20)
GAE_Trainer.train_iters(tr, dev, te, model_file="./tmp.pkl")

# grd_cls = nn.Linear(300, 2)
# grd = GraphRumorDetection(w2v, gat, grd_cls, batch_size=20, grad_accum_cnt=1)
#
# grd.train_iters(tr, dev, te,
#             valid_every=100, max_epochs=100, lr_discount=1.0,
#             log_dir="../logs/", log_suffix="_RumorDetection", model_file="../saved/GAE_RDM_CN.pkl")
